#include "ffmpeg.h"
#include "v4l2.h"
AVCodecContext *global_codec_ctx = nullptr; //H.264编码上下文(编解码)
AVFormatContext *global_fmt_ctx = nullptr;  //格式上下文 mp4容器
struct SwsContext *global_sws_ctx = nullptr; //格式转换上下文
AVFrame *global_frame = nullptr; //视频帧缓冲区
AVPacket *global_pkt = nullptr; //编码后的数据包
AVStream *global_stream = nullptr;

int ffmpeg_init(const char* filename){

    av_log_set_level(AV_LOG_DEBUG); // 启用 FFmpeg 调试日志

    //search encoder H264
    const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        std::cerr << "未找到 H.264 编码器" << std::endl;
        return -1;
    }

    // 编码上下文
    global_codec_ctx = avcodec_alloc_context3(codec);
    if(!global_codec_ctx){
        std::cerr << "无法分配H.264编码上下文" << std::endl;
        return -1;
    }
    global_codec_ctx->bit_rate = 300000; // 比特率
    global_codec_ctx->width = WIDTH;
    global_codec_ctx->height = HEIGHT;
    global_codec_ctx->time_base = (AVRational){1, FPS};
    global_codec_ctx->framerate = (AVRational){FPS, 1};
    global_codec_ctx->gop_size = 12;
    global_codec_ctx->max_b_frames = 1;
    global_codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
    //设置
    av_opt_set(global_codec_ctx->priv_data, "preset", "superfast", 0);
    av_opt_set(global_codec_ctx->priv_data, "tune", "zerolatency", 0);
    av_opt_set(global_codec_ctx->priv_data, "threads", "4", 0); // 利用4核CPU

    //打开编码器
    if (avcodec_open2(global_codec_ctx, codec, NULL) < 0) {
        std::cerr << "无法打开H.264编码器" << std::endl;
        return -1;
    }

    //格式上下文(创建MP4格式)
    if (avformat_alloc_output_context2(&global_fmt_ctx, NULL, "mp4", filename) < 0) {
        std::cerr << "无法创建输出上下文" << std::endl;
        return -1;
    }

    // 添加视频流(H264)
    global_stream = avformat_new_stream(global_fmt_ctx, codec);
    if (!global_stream) {
        std::cerr << "无法创建视频流" << std::endl;
        return -1;
    }

    //编码上下文参数传递给流参数
    avcodec_parameters_from_context(global_stream->codecpar, global_codec_ctx);

    // 打开输出文件
    if (avio_open(&global_fmt_ctx->pb, filename, AVIO_FLAG_WRITE) < 0) {
        std::cerr << "无法打开输出文件: " <<filename <<" - " << strerror(errno) << std::endl;
        return -1;
    }

    // 写入文件头
    if (avformat_write_header(global_fmt_ctx, NULL) < 0) {
        std::cerr << "无法写入文件头" << std::endl;
        return -1;
    }

    // --- 格式转换初始化 --- // Convert UYVY to YUV420P for MP4
    global_sws_ctx = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_UYVY422,
                                    WIDTH, HEIGHT, AV_PIX_FMT_YUV420P,
                                    SWS_BILINEAR, NULL, NULL, NULL);
    if (!global_sws_ctx) {
        std::cerr << "无法初始化格式转换上下文" << std::endl;
        return -1;
    }


    // --- 帧分配 ---
    global_frame = av_frame_alloc();
    if (!global_frame) {
        std::cerr << "无法分配帧" << std::endl;
        return -1;
    }
    global_frame->format = global_codec_ctx->pix_fmt;
    global_frame->width = global_codec_ctx->width;
    global_frame->height = global_codec_ctx->height;
    if (av_frame_get_buffer(global_frame, 32) < 0) {  //32字节对齐
        std::cerr << "无法分配帧缓冲区" << std::endl;
        return -1;
    }

    // --- 数据包分配 ---
    global_pkt = av_packet_alloc();
    if (!global_pkt) {
        std::cerr << "无法分配数据包" << std::endl;
        return -1;
    }

    return 0;
}

//编码、写入MP4
int encode_and_write_frame(uint8_t* frame_data, int frame_count) {
    // 将 UYVY 数据转换为 YUV420P
    uint8_t *src_data[1] = { frame_data };
    int src_linesize[1] = { WIDTH * 2 };
    
    //视频像素格式和分辨率的转换
    //sws_getContext返回的参数  input    input       0       高度            output           output
    sws_scale(global_sws_ctx, src_data, src_linesize, 0, HEIGHT, global_frame->data, global_frame->linesize);
    // 设置帧时间戳
    global_frame->pts = av_rescale_q(frame_count, (AVRational){1, FPS}, global_codec_ctx->time_base);

    // 发送帧(global_frame)到编码器
    if (avcodec_send_frame(global_codec_ctx, global_frame) < 0) {
        std::cerr << "发送帧到编码器失败" << std::endl;
        return -1;
    }

    // 接收编码后的数据包
    while (1) {
        int ret = avcodec_receive_packet(global_codec_ctx, global_pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            break;
        if (ret < 0) {
            std::cerr << "编码过程中出错: " << ret << std::endl;
            break;
        }
        global_pkt->stream_index = global_stream->index;

        //转换时间戳
        av_packet_rescale_ts(global_pkt, global_codec_ctx->time_base, global_stream->time_base);
        // 将数据包写入输出格式上下文
        if (av_interleaved_write_frame(global_fmt_ctx, global_pkt) < 0) {  //写入mp4文件
            std::cerr << "写入数据包失败" << std::endl;
            return -1;
        }
        av_packet_unref(global_pkt); //清理AVPacket中的所有空间数据
    }
    return 0;
}

//拍照
int save_snapshot(uint8_t * frame_data, const char* filename){

    // Initialize JPEG encoder 
    const AVCodec *jpeg_codec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (!jpeg_codec) {
        std::cerr << "未找到 JPEG 编码器" << std::endl;
        return -1;
    }

    // 编码器上下文
    AVCodecContext *jpeg_ctx = avcodec_alloc_context3(jpeg_codec);
    if(!jpeg_ctx){
        std::cerr << "无法分配JPEG编码上下文" << std::endl;
        return -1;
    }
    jpeg_ctx->width = WIDTH;
    jpeg_ctx->height = HEIGHT;
    jpeg_ctx->pix_fmt = AV_PIX_FMT_YUVJ420P;
    jpeg_ctx->time_base = (AVRational){1, FPS};
    av_opt_set_int(jpeg_ctx->priv_data, "q", 2, 0);

    //打开编码器
    if (avcodec_open2(jpeg_ctx, jpeg_codec, NULL) < 0) {
        std::cerr << "无法打开JPEG编码器" << std::endl;
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // Convert UYVY to YUV420P
    SwsContext *sws_ctx = sws_getContext(WIDTH, HEIGHT, AV_PIX_FMT_UYVY422,
                        WIDTH, HEIGHT, AV_PIX_FMT_YUV420P,
                        SWS_BILINEAR, NULL ,NULL ,NULL);
    
    if (!sws_ctx) {
        std::cerr << "无法初始化拍照转换上下文" << std::endl;
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // --- 帧分配 ---
    AVFrame *frame = av_frame_alloc();
    if (!frame) {
        std::cerr << "无法分配帧" << std::endl;
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }
    frame->format = AV_PIX_FMT_YUV420P;
    frame->width = WIDTH;
    frame->height = HEIGHT;
    if (av_frame_get_buffer(frame, 32) < 0) {  //设置好宽度、高度、格式，自动为AVFrame分配空间 32字节对齐
        std::cerr << "无法分配帧缓冲区" << std::endl;
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    uint8_t *src_data[1] = {frame_data};
    int src_linesize[1] = {WIDTH * 2};
    //像素格式分辨率转换     in           in                    out             out
    sws_scale(sws_ctx, src_data, src_linesize ,0 ,HEIGHT, frame->data, frame->linesize);

    //Encode frame
    AVPacket *pkt = av_packet_alloc();
    if (!pkt) {
        std::cerr << "无法分配数据包" << std::endl;
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // 发送帧(frame)到编码器
    if (avcodec_send_frame(jpeg_ctx, frame) < 0) {
        std::cerr << "发送帧到编码器失败" << std::endl;
        av_packet_free(&pkt);
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    // 接收编码后的数据包
    if(avcodec_receive_packet(jpeg_ctx, pkt) < 0){
        std::cerr << "Failed to receive JPEG packet" << std::endl;
        av_packet_free(&pkt);
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }

    //write to file
    FILE *file = fopen(filename, "wb");
    if(!file){
        std::cerr << "Failed to open snapshot file: " << strerror(errno) << std::endl;
        av_packet_free(&pkt);
        av_frame_free(&frame);
        sws_freeContext(sws_ctx);
        avcodec_free_context(&jpeg_ctx);
        return -1;
    }
    std::cout<< "open snapshot file ok" <<std::endl;
    //压缩编码的数据  被写入的每个元素的大小，以字节为单位   数据的大小
    fwrite(pkt->data, 1, pkt->size, file);
    fclose(file);

    // Clean up
    av_packet_free(&pkt);
    av_frame_free(&frame);
    sws_freeContext(sws_ctx);
    avcodec_free_context(&jpeg_ctx);
    return 0;
    

}

void ffmpeg_cleanup(void) {
    // Flush encoder
    avcodec_send_frame(global_codec_ctx, nullptr);
    while (1) {
        int ret = avcodec_receive_packet(global_codec_ctx, global_pkt);
        if (ret == AVERROR_EOF || ret == AVERROR(EAGAIN))
            break;
        if (ret < 0) {
            std::cerr << "Failed to flush encoder" << std::endl;
            break;
        }
        global_pkt->stream_index = global_stream->index;
        //调整时间戳
        av_packet_rescale_ts(global_pkt, global_codec_ctx->time_base, global_stream->time_base);
        if (av_interleaved_write_frame(global_fmt_ctx, global_pkt) < 0) { //将数据包写入文件mp4
            std::cerr << "Failed to write flushed packet" << std::endl;
        }
        av_packet_unref(global_pkt); //释放数据包的引用，避免内存泄漏。
    }

    // Write file trailer
    if (global_fmt_ctx) {
        av_write_trailer(global_fmt_ctx);
    }

    // Clean up resources
    if (global_pkt) av_packet_free(&global_pkt);
    if (global_frame) av_frame_free(&global_frame);
    if (global_sws_ctx) sws_freeContext(global_sws_ctx);
    if (global_codec_ctx) avcodec_free_context(&global_codec_ctx);
    if (global_fmt_ctx) {
        avio_closep(&global_fmt_ctx->pb); // 关闭流并释放资源
        avformat_free_context(global_fmt_ctx);
    }
}



